#include "memalloc.h"
#include <pthread.h>
#include "utils.h"

struct memory_header {
    char reserved[16];
    int size;
    struct memory_header* prev;
    struct memory_header* next;
};

struct memory_buffer {
    int count;
    int size;
    unsigned char* buffer;
    struct memory_header* idle;
};

struct memory_buffer membuf_big;        // 大内存缓冲区
struct memory_buffer membuf_small[5];   // 小内存缓冲区，会预先分配一定数量的内存块
int membuf_sizes[] = { 128, 512, 1024, 4096, 1024*8 };
int membuf_counts[] = { 64, 64, 64, 48, 24 };
pthread_mutex_t malloc_lock;

void init_memory_buffer(struct memory_buffer* membuf, int size, int count)
{
    struct memory_header *p, *prev;
    unsigned char* buffer;
    int i, sz = size + sizeof(struct memory_header);

    buffer = (unsigned char*)malloc(sz * count);
    memset(buffer, 0, sz * count);

    i = 0;
    prev = NULL;
    for (i = 0; i < count; i++) {
        p = (struct memory_header*)(buffer + i * sz);
        p->size = size;
        p->prev = prev;
        if (prev)
            prev->next = p;
        prev = p;
    }
    
    membuf->count = count;
    membuf->buffer = buffer;
    membuf->idle = (struct memory_header*)buffer;
}

void free_all(struct memory_header *first, struct memory_header *last, struct memory_header *p)
{
    while (p) {
        if (p < first || p > last) {
            struct memory_header *p1 = p;
            p = p->next;
            free(p1);
        }
        else
            p = p->next;
    }
}

void free_memory_buffer(struct memory_buffer* membuf)
{
    struct memory_header *first, *last;

    if (membuf->count > 0) {
        int sz = membuf->size + sizeof(struct memory_header);

        first = (struct memory_header*)membuf->buffer;
        last = (struct memory_header*)(membuf->buffer + (membuf->count - 1) * sz);
    }
    else {
        first = NULL;
        last = NULL;
    }
    free_all(first, last, membuf->idle);

    if (membuf->buffer) {
        free(membuf->buffer);
    }
}

void malloc_init()
{
    memset(&membuf_small, 0, sizeof(membuf_small));
    memset(&membuf_big, 0, sizeof(membuf_big));

    pthread_mutex_init(&malloc_lock, 0);
}

void malloc_uninit()
{
    int i;

    pthread_mutex_lock(&malloc_lock);
    
    for (i = 0; i < 5; i++) {
        free_memory_buffer(&membuf_small[i]);
    }
    free_memory_buffer(&membuf_big);

    pthread_mutex_unlock(&malloc_lock);

    pthread_mutex_destroy(&malloc_lock);
}

void* mem_alloc(size_t size)
{
    struct memory_buffer* membuf;
    struct memory_header* mptr = NULL;

    pthread_mutex_lock(&malloc_lock);

    if (size <= membuf_sizes[4]) {
        int index = 4;

        if (size <= membuf_sizes[0]) index = 0;
        else if (size <= membuf_sizes[1]) index = 1;
        else if (size <= membuf_sizes[2]) index = 2;
        else if (size <= membuf_sizes[3]) index = 3;

        membuf = &membuf_small[index];
        size = (size_t)membuf_sizes[index];
        if (membuf->count == 0)
            init_memory_buffer(membuf, (int)size, membuf_counts[index]);

        if (membuf->idle) {
            // 有空闲内存
            mptr = membuf->idle;
            membuf->idle = mptr->next;
            if (membuf->idle)
                membuf->idle->prev = NULL;
        }
    }
    else {
        membuf = &membuf_big;
        if (membuf->idle) {
            // 选择合适大小的空闲内存
            struct memory_header *p1, *p2 = NULL;
            int delta_size = 0;

            p1 = membuf->idle;
    // dlog(LOG_INFO, "  || idle: %s", print_buffers(p1));
            while (p1) {
                if (p1->size == (int)size) {
                    p2 = p1;
                    delta_size = 0;
                    break;
                }
                if (p1->size > (int)size) {
                    int nsize = p1->size - (int)size;
                    if (!p2 || delta_size > nsize) {
                        p2 = p1;
                        delta_size = nsize;
                    }
                }
                p1 = p1->next;
            }

            if (p2) {
                if (delta_size == 0) {
                    mptr = p2;
                } else {
                    if (size < 1024 * 64) {
                        if (delta_size < 1024 * 4)
                            mptr = p2;
                    }
                    else if (size < 1024 * 128) {
                        if (delta_size < 1024 * 8)
                            mptr = p2;
                    }
                    else if (delta_size < 1024 * 32)
                        mptr = p2;
                }

                if (mptr) {
                    if (mptr->prev) {
                        mptr->prev->next = mptr->next;
                        if (mptr->next)
                            mptr->next->prev = mptr->prev;
                    }
                    else {
                        membuf->idle = mptr->next;
                        if (membuf->idle)
                            membuf->idle->prev = NULL;
                    }
                }
            }
        }
    }

    if (!mptr) {
        // 没有空闲内存，分配新的内存
        mptr = (struct memory_header*)malloc(sizeof(struct memory_header) + size);
        if (!mptr) {
            pthread_mutex_unlock(&malloc_lock);
            return NULL;
        }
        mptr->size = (int)size;
    }
    pthread_mutex_unlock(&malloc_lock);

    return (void*)((unsigned char*)mptr + sizeof(struct memory_header));
}

void* mem_zalloc(size_t size)
{
    void* p = mem_alloc(size);
	if (p)
		memset(p, 0, size);
    return p;
}

void mem_free(void* ptr)
{
    struct memory_buffer* membuf;
    struct memory_header* mptr;
    if (!ptr)
        return;

    pthread_mutex_lock(&malloc_lock);

    mptr = (struct memory_header*)((unsigned char*)ptr - sizeof(struct memory_header));
    if (mptr->size <= membuf_sizes[4]) {
        int index = 4;

        if (mptr->size <= membuf_sizes[0]) index = 0;
        else if (mptr->size <= membuf_sizes[1]) index = 1;
        else if (mptr->size <= membuf_sizes[2]) index = 2;
        else if (mptr->size <= membuf_sizes[3]) index = 3;

        membuf = &membuf_small[index];
    }
    else {
        membuf = &membuf_big;
    }

    mptr->prev = NULL;
    mptr->next = membuf->idle;
    if (membuf->idle)
        membuf->idle->prev = mptr;
    membuf->idle = mptr;

    pthread_mutex_unlock(&malloc_lock);
}

/*  重新分配一个更大的内存
        ptr         原内存指针，为空则直接分配新的内存
        newsize     分配新内存的大小，小于等于旧内存时，直接返回旧内存指针
    分配新内存后，将旧内存的内容复制到新内存中，并释放旧内存
*/
void* mem_realloc(void* ptr, size_t newsize)
{
    struct memory_header* mptr;
    void* newptr;
    if (!ptr)
        return mem_alloc(newsize);

    mptr = (struct memory_header*)((unsigned char*)ptr - sizeof(struct memory_header));
    if (newsize <= mptr->size)
        return ptr;

    newptr = mem_alloc(newsize);  // 分配新内存
    if (newptr) {
        memcpy(newptr, ptr, mptr->size);    // 复制旧内存的内容到新内存中
    }
    mem_free(ptr); // 释放旧内存
    return newptr;
}

char* mem_strdup(const char* s)
{
    int len;
    char* ns;

    if (!s)
        return 0;
    
    len = (int)strlen(s);
    ns = (char *)mem_alloc(len + 1);
    strcpy(ns, s);

    return ns;
}
